import os
import cv2
import numpy as np
import shutil
from insightface.app import FaceAnalysis
from sklearn.metrics.pairwise import cosine_similarity

class FaceClassifier:
    def __init__(self, source_path, output_path, similarity_threshold=0.4):
        """
        初始化人脸分类器
        
        参数:
        source_path: 源图片目录路径
        output_path: 分类结果输出目录
        similarity_threshold: 人脸相似度阈值（默认0.6）
        """
        self.source_path = source_path
        self.output_path = output_path
        self.similarity_threshold = similarity_threshold
        
        # 初始化人脸分析模型
        self.app = FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))
        
        # 存储已知人物信息
        self.known_persons = {}
        self.current_person_id = 0
        
        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)

    def process_images(self):
        """处理源目录中的所有图片"""
        # 获取所有支持的图片文件
        image_exts = ('.jpg', '.jpeg', '.png', '.bmp', '.tif')
        image_files = [f for f in os.listdir(self.source_path) 
                      if f.lower().endswith(image_exts)]
        
        print(f"发现 {len(image_files)} 张图片需要处理")
        
        # 处理每张图片
        for i, img_file in enumerate(image_files):
            img_path = os.path.join(self.source_path, img_file)
            print(f"\n处理图片 {i+1}/{len(image_files)}: {img_file}")
            
            # 读取并处理图片
            img = cv2.imread(img_path)
            if img is None:
                print(f"  警告：无法读取图片 {img_file}，跳过")
                continue
            
            # 检测人脸
            faces = self.app.get(img)
            
            if not faces:
                print(f"  未检测到人脸，跳过")
                self._save_unclassified(img, img_file)
                continue
                
            # 处理每张检测到的人脸
            for j, face in enumerate(faces):
                face_embedding = face.embedding
                bbox = face.bbox.astype(int)
                
                # 裁剪人脸区域
                face_img = img[bbox[1]:bbox[3], bbox[0]:bbox[2]]
                
                # 识别人物
                person_id, similarity = self._recognize_person(face_embedding)
                
                if person_id:
                    print(f"  人脸 {j+1}: 识别为 {person_id} (相似度: {similarity:.4f})")
                    self._save_person_face(person_id, face_img, img_file)
                else:
                    # 注册新人物
                    self.current_person_id += 1
                    new_person = f"Person_{self.current_person_id}"
                    self.known_persons[new_person] = face_embedding
                    print(f"  人脸 {j+1}: 注册新人物 {new_person} (相似度: {similarity:.4f})")
                    self._save_person_face(new_person, face_img, img_file)
        
        print("\n处理完成!")
        print(f"共识别出 {len(self.known_persons)} 个不同人物")
        print(f"分类结果保存在: {self.output_path}")

    def _recognize_person(self, embedding):
        """识别人物并返回人物ID和相似度"""
        max_similarity = 0
        recognized_person = None
        
        for person_id, known_embedding in self.known_persons.items():
            # 计算余弦相似度
            sim = cosine_similarity([embedding], [known_embedding])[0][0]
            
            if sim > max_similarity:
                max_similarity = sim
                recognized_person = person_id
        
        # 检查是否达到相似度阈值
        if max_similarity >= self.similarity_threshold:
            return recognized_person, max_similarity
        else:
            return None, max_similarity

    def _save_person_face(self, person_id, face_img, original_filename):
        """保存人脸图片到对应人物目录"""
        person_dir = os.path.join(self.output_path, person_id)
        os.makedirs(person_dir, exist_ok=True)
        
        # 生成唯一文件名
        base_name = os.path.splitext(original_filename)[0]
        face_filename = f"{base_name}_{os.urandom(4).hex()}.jpg"
        face_path = os.path.join(person_dir, face_filename)
        
        # 保存人脸图片
        cv2.imwrite(face_path, face_img)

    def _save_unclassified(self, img, original_filename):
        """保存未分类的图片"""
        unclassified_dir = os.path.join(self.output_path, "unclassified")
        os.makedirs(unclassified_dir, exist_ok=True)
        cv2.imwrite(os.path.join(unclassified_dir, original_filename), img)

if __name__ == "__main__":
    # 配置路径参数
    SOURCE_DIR = "input_images/"  # 替换为你的图片目录路径
    OUTPUT_DIR = "output/"       # 替换为输出目录路径
    
    # 创建分类器并处理图片
    classifier = FaceClassifier(source_path=SOURCE_DIR, output_path=OUTPUT_DIR)
    classifier.process_images()
